import cv2
import numpy as np
import os

# 全局变量，用于保存鼠标点击的像素坐标
clicked_point = None
def get_image_paths(directory):
    # 定義常見圖片格式
    # image_extensions = (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".webp")
    image_extensions = (".raw")
    image_paths = []

    # 遍歷目錄及其子目錄
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(image_extensions):  # 判斷是否是圖片
                image_paths.append(os.path.join(root, file))
    
    # 對結果進行排序
    image_paths.sort()  # 按完整路徑排序
    return image_paths

def mouse_callback(event, x, y, flags, param):
    """鼠标点击回调函数"""
    global clicked_point
    if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
        clicked_point = (x, y)
        print(f"选中的像素点坐标: {clicked_point}")

def get_xyz_from_pixel(pixel: tuple, depth_map: np.ndarray, K: np.ndarray, D: np.ndarray) -> tuple:
    """根据像素点坐标获取点云空间内的 (x, y, z)"""

    # 获取的是 x y
    u, v = pixel  # x,y
    print(f'ori {u, v}')
    # uv 行 列
    # u = int( u / 768 * 480)
    # v = int( v / 1024 * 640)
    print(f'cvt {u, v}')


    # 获取深度值（单位：米）
    # z = depth_map[u, v] / 1000.0  # 假设深度图单位是毫米，转换为米
    z = depth_map[v][u] / 1000.0  # 假设深度图单位是毫米，转换为米
    print(f'z {z}')
    if z <= 0 or z > 20.0:  # 检查深度值是否有效
    

        neighbors = [
            (u - 1, v - 1), (u, v - 1), (u + 1, v - 1),  # 上方三个点
            (u - 1, v),                 (u + 1, v),      # 左右两个点
            (u - 1, v + 1), (u, v + 1), (u + 1, v + 1)   # 下方三个点
        ]
        
        for nu, nv in neighbors:
            if 0 <= nv < depth_map.shape[0] and 0 <= nu < depth_map.shape[1]:  # 检查是否在图像范围内
                neighbor_z = depth_map[nv][nu] / 1000.0  # 获取深度值
                print(f"邻居像素点 ({nu}, {nv}) 的深度值: {neighbor_z} 米")
            else:
                print(f"邻居像素点 ({nu}, {nv}) 超出图像范围")
                      
        raise ValueError(f"无效的深度值: {z} m")
        
    
    # 使用鱼眼模型进行反投影
    pts = np.array([[[u, v]]], dtype=np.float32)  # 像素点坐标需要是 (1, 1, 2) 的形状
    undistorted_pts = cv2.fisheye.undistortPoints(pts, K, D)  # 去畸变并反投影
    x = undistorted_pts[0][0][0] * z
    y = undistorted_pts[0][0][1] * z
    return x, y, z

def load_depth_map(depth_path: str, depth_height: int = 480, depth_width: int = 640) -> np.ndarray:
    with open(depth_path, 'rb') as f:
        depth_data = np.fromfile(f, dtype=np.uint16)

    depth_map = depth_data.reshape((depth_height, depth_width))
    return depth_map


def compute_depth_pseudo_color(depth_image: np.ndarray, max_depth: float = 5.0) -> np.ndarray:
    """将深度图转为伪彩色图像"""
    depth_normalized = np.clip(depth_image / (max_depth * 1000), 0, 1)
    depth_colored = cv2.applyColorMap((depth_normalized * 255).astype(np.uint8), cv2.COLORMAP_JET)
    return depth_colored



def get_depth_mask_path(main_path, choice, function, image_paths):
    choice = choice - 1
    depth_path = image_paths[choice]

    depth_map = load_depth_map(depth_path)
    resized_depth_map = cv2.resize(depth_map, (1024, 768), interpolation=cv2.INTER_NEAREST)


    file_name_with_ext = os.path.basename(depth_path)  # 提取文件名（带扩展名）
    file_id = os.path.splitext(file_name_with_ext)[0]  # 去掉扩展名，获取文件 ID

    new_name = None
    if function == 0:
        new_name = os.path.join(main_path, "left_reproject_undistort_fisheye", file_id+".jpg")
    if function == 1:
        new_name = os.path.join(main_path, "mask_index", file_id+"_"+str(choice+1)+".jpg")

    show_img =  cv2.imread(new_name)

    return resized_depth_map, show_img



def main():
    global clicked_point  # 声明 clicked_point 为全局变量

    K = np.array([[380.514984, 0, 512.130066],
                  [0, 380.059082, 332.302795],
                  [0, 0, 1]]) 

    D = np.array([0.0422173887, -0.000108877211, -0.000979754375, -0.000220517832])

    T_v_cam = np.array([
        [0.00274878, -0.141637, 0.989915, 0.466449],
        [-0.999923, 0.0115661, 0.00443145, 0.0242619],
        [-0.0120771, -0.989851, -0.141595, 0.111014],
        [0.0, 0.0, 0.0, 1.0]
    ])

    main_path = '/media/ai/0bea6433-71ce-4bf1-a689-3b0348c3c57b/vslam/0423/965/results'
    raw_path = os.path.join(main_path, "depth")
    image_paths = get_image_paths(raw_path)

    print(f'all path.size() {image_paths.__len__()}')
 

    while True:
        print("提示：输入图片序号查看图片，按ESC键退出程序。")
        try:
            choice = input("请输入要查看的图片序号：")
            print(choice)
            if choice.strip().lower() == 'esc':  # 去除首尾空格并转换为小写

                print("程序已退出。")
                break
            choices = choice.split()
            if len(choices) == 1:  # 如果只输入一个值
                index = int(choices[0])
                if 0 <= index < len(image_paths):
                    pass
                else:
                    print("无效的序号，请重新输入。")
            elif len(choices) == 2:  # 如果输入两个值
                index = int(choices[0])
                function = int(choices[1])
                depth_map, show_img = get_depth_mask_path(main_path, index, function, image_paths)
                # 显示 RGB 图像并绑定鼠标回调
                cv2.namedWindow("RGB Image")
                cv2.setMouseCallback("RGB Image", mouse_callback)
                display_image = show_img.copy()
                while True:
                    # 显示 RGB 图像

                    if clicked_point is not None:
                        try:
                            # 获取点击点的 (x, y, z)
                            x, y, z = get_xyz_from_pixel(clicked_point, depth_map, K, D)
                            # x, y, z = get_xyz_from_pixel(clicked_point, resized_depth_map, K, D)
                            print(f"像素点 {clicked_point} 对应的相机坐标: x={x:.3f}, y={y:.3f}, z={z:.3f}")
                            p_cam = np.array([x, y, z, 1.0])  # 示例点 (x, y, z, 1)

                            p_v = np.dot(T_v_cam, p_cam) # 

                            new_x, new_y, new_z, _ = p_v

                            print(f"像素点 {clicked_point} 对应的世界坐标: x={new_x:.3f}, y={new_y:.3f}, z={new_z:.3f}")


                            
                            # 在图像上显示深度值和点云坐标
                            # cv2.putText(display_image, f"Depth: {z:.3f}m", (clicked_point[0], clicked_point[1] - 10),
                            #             cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
                            # cv2.putText(display_image, f"X: {x:.3f}, Y: {y:.3f}, Z: {z:.3f}", (10, 30),
                            #             cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                            clicked_point = None  # 重置点击点
                        except ValueError as e:
                            print(e)
                            clicked_point = None  # 重置点击点

                    cv2.imshow("RGB Image", display_image)
                    key = cv2.waitKey(1) & 0xFF

                    if key == ord('q'):  # 按下 'q' 键退出
                        break

                cv2.destroyAllWindows()


        except ValueError:
            print("请输入有效的数字或ESC退出。")


 




if __name__ == "__main__":
    main()
